"""React智能体-Mysql工具"""

import json
import time
from typing import Any, Dict, List, Tuple
from json_repair import repair_json

try:
    import mysql.connector
    from mysql.connector import Error
    MYSQL_AVAILABLE = True
except ImportError:
    MYSQL_AVAILABLE = False
    mysql = None
    Error = Exception

from typing import Type, Optional
from pydantic import BaseModel, Field
from langchain_core.callbacks import (
    AsyncCallbackManagerForToolRun,
    CallbackManagerForToolRun,
)

from .tool_result import ToolResult

import logging
logger = logging.getLogger(__name__)


class MySQLInput(BaseModel):
    query: str = Field(
        description="""MySQL 数据库查询格式：' 操作 [参数...]'。操作如下：
- list_tables：列出数据库中的所有表
- describe <表名>：获取表结构
- get <表名> [where 子句]：从表中获取记录
- set <表名> <json 数据 >：插入新记录
- update <表名> <json 数据 > where < 条件 >：更新现有记录
- delete <表名> where < 条件 >：删除记录
- search <术语>：跨所有表搜索
- stats：获取数据库统计信息
- create_table <表名> < 列 >：创建新表
- sql <查询语句>：执行自定义 SQL 查询
""")

class MySQLConnection:

    def __init__(self, host: str, database: str, user: str, password: str, port: int = 3306):
        self.host = host
        self.database = database
        self.user = user
        self.password = password
        self.port = port
        self.connection = None
        self._table_cache = {}
        self._cache_timestamp = 0
        self._cache_ttl = 300
    
    def connect(self):
        if not MYSQL_AVAILABLE:
            raise ImportError("MySQL connector not available")
        try:
            self.connection = mysql.connector.connect(
                host=self.host,
                database=self.database,
                user=self.user,
                password=self.password,
                port=self.port,
                autocommit=True,
                charset='utf8mb4',
                use_unicode=True
            )
            logger.info(f"MySQL连接成功: {self.database}")
            return True
        except Error as e:
            logger.error(f"MySQL连接失败: {e}")
            return False
    
    def disconnect(self):
        if self.connection and self.connection.is_connected():
            self.connection.close()
            logger.info("MySQL连接已关闭")
    
    def is_connected(self) -> bool:
        return self.connection and self.connection.is_connected()
    
    def ensure_connection(self):
        if not self.is_connected():
            self.connect()
    
    def execute_query(self, query: str, params: Tuple = None) -> Tuple[bool, Any, str]:
        try:
            self.ensure_connection()
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute(query, params or ())
            
            if query.strip().upper().startswith(('SELECT', 'SHOW', 'DESCRIBE', 'EXPLAIN')):
                result = cursor.fetchall()
            else:
                result = cursor.rowcount
            
            cursor.close()
            return True, result, ""
            
        except Error as e:
            logger.error(f"MySQL query error: {e}")
            return False, None, str(e)
    
    def get_table_names(self) -> List[str]:
        success, result, error = self.execute_query("SHOW TABLES")
        if success and result:
            # MySQL returns table names in format: {'Tables_in_dbname': 'table_name'}
            table_key = f"Tables_in_{self.database}"
            return [row[table_key] for row in result]
        return []
    
    def get_table_schema(self, table_name: str) -> Dict[str, Any]:
        cache_key = f"{table_name}_schema"
        if (cache_key in self._table_cache and 
            time.time() - self._cache_timestamp < self._cache_ttl):
            return self._table_cache[cache_key]
        
        success, result, error = self.execute_query(f"DESCRIBE {table_name}")
        if success and result:
            schema = {
                "columns": [],
                "primary_keys": [],
                "column_types": {}
            }
            
            for column in result:
                col_name = column['Field']
                col_type = column['Type']
                is_primary = column['Key'] == 'PRI'
                
                schema["columns"].append(col_name)
                schema["column_types"][col_name] = col_type
                if is_primary:
                    schema["primary_keys"].append(col_name)

            self._table_cache[cache_key] = schema
            self._cache_timestamp = time.time()
            return schema
        
        return {"columns": [], "primary_keys": [], "column_types": {}}
    
    def table_exists(self, table_name: str) -> bool:
        return table_name in self.get_table_names()


class MySQLDatabaseTool(BaseModel):
    name: str = 'mysql_database'
    description: str = """在MySQL数据库中查询和管理数据。

操作与语法：
• list_tables - 列出所有表
  用法：list_tables

• describe <table> - 显示表结构
  用法：describe users

• get <table> [conditions] - 检索数据
  用法：get users
  用法：get users id=1
  用法：get users * 10（限制10条记录）

• set <table> <json_data> - 插入新记录
  用法：set users {"name": "John", "email": "john@email.com", "age": 25}
  ⚠️ 重要提示：数据必须是带双引号的有效JSON对象

• update <table> <json_data> WHERE <condition> - 更新记录
  用法：update users {"name": "Jane"} WHERE id=1

• delete <table> WHERE <condition> - 删除记录
  用法：delete users WHERE id=1

• search <table> <column> <value> - 搜索记录
  用法：search users name John

• stats - 显示数据库统计信息
  用法：stats

• create_table <table_name> (<columns>) - 创建表
  用法：create_table users (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100), age INT)

• sql <query> - 执行自定义SQL
  用法：sql SELECT * FROM users WHERE age > 18

• help - 显示详细使用信息
  用法：help

常见错误：
- JSON解析错误 → 检查“set”命令中的JSON格式
- 表不存在 → 先使用“list_tables”
- 语法错误 → 参考上面的示例"""

    args_schema: Type[BaseModel] = MySQLInput.model_json_schema()
    return_direct: bool = True

    def __init__(self, host: str, database: str, user: str, password: str, port: int = 3306):
        super().__init__()
        if not MYSQL_AVAILABLE:
            raise ImportError("MySQL 连接器不可用，请先安装 mysql-connector-python。")
        self.mysql = MySQLConnection(host, database, user, password, port)
        self._connect_to_database()

    def _connect_to_database(self):
        if not self.mysql.connect():
            raise Exception("MySQL 连接失败")

    def _run(self, query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> ToolResult:
        raise NotImplementedError("不支持同步调用. ")

    async def _arun(self, query: str,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None
    ) -> ToolResult:
        """执行Mysql数据库操作"""
        try:
            parts = query.strip().split()
            if not parts:
                return ToolResult(success=False, data=None, error="查询为空.")
            
            operation = parts[0].lower()
            if operation == "list_tables":
                return await self._list_tables_operation()
            elif operation == "describe" or operation == "desc":
                return await self._describe_table_operation(parts[1:])
            elif operation == "get":
                return await self._get_operation(parts[1:])
            elif operation == "set" or operation == "insert":
                return await self._set_operation(parts[1:])
            elif operation == "update":
                return await self._update_operation(parts[1:])
            elif operation == "delete":
                return await self._delete_operation(parts[1:])
            elif operation == "search":
                return await self._search_operation(parts[1:])
            elif operation == "stats":
                return await self._stats_operation()
            elif operation == "create_table":
                return await self._create_table_operation_improved(query)
            elif operation == "custom_query" or operation == "sql":
                return await self._custom_query_operation(parts[1:])
            elif operation == "help":
                return await self._help_operation()
            else:
                return ToolResult(success=False, data=None,
                    error=f"未知操作：{operation}。支持的操作包括：list_tables、describe、get、set、update、delete、search、stats、create_table、sql、help。使用“help”可获取详细的使用信息。"
                )
        except Exception as e:
            logger.error(f"数据库操作失败：{e}")
            return ToolResult(success=False, data=None, error=f"数据库操作失败：{str(e)}")
    
    async def _list_tables_operation(self) -> ToolResult:
        tables = self.mysql.get_table_names()
        return ToolResult(
            success=True,
            data={"tables": tables, "count": len(tables)},
            metadata={"operation": "list_tables", "database": self.mysql.database}
        )
    
    async def _describe_table_operation(self, args: List[str]) -> ToolResult:
        if not args:
            return ToolResult(success=False,data=None, error="操作描述需要表名信息")
        
        table_name = args[0]
        if not self.mysql.table_exists(table_name):
            return ToolResult(success=False, data=None, error=f"表 '{table_name}' 不存在")
        
        schema = self.mysql.get_table_schema(table_name)
        success, result, error = self.mysql.execute_query(f"DESCRIBE {table_name}")
        
        if not success:
            return ToolResult(success=False, data=None, error=f"获取表的结构信息失败: {error}")
        
        return ToolResult(
            success=True,
            data={
                "table": table_name,
                "schema": result,
                "columns": schema["columns"],
                "primary_keys": schema["primary_keys"],
                "column_types": schema["column_types"]
            },
            metadata={"operation": "describe", "table": table_name}
        )
    
    async def _get_operation(self, args: List[str]) -> ToolResult:
        if not args:
            return ToolResult(success=False, data=None, error="操作描述需要表名信息")
        
        table_name = args[0]
        if not self.mysql.table_exists(table_name):
            return ToolResult( success=False, data=None, error=f"表 '{table_name}' 不存在")

        if len(args) == 1:
            query = f"SELECT * FROM {table_name} LIMIT 100"
        elif len(args) == 2 and args[1] == "*":
            query = f"SELECT * FROM {table_name}"
        elif len(args) >= 3 and args[1] == "*":
            try:
                limit = int(args[2])
                query = f"SELECT * FROM {table_name} LIMIT {limit}"
            except ValueError:
                return ToolResult(
                    success=False,
                    data=None,
                    error="Invalid limit value"
                )
        else:
            where_clause = " ".join(args[1:])
            if where_clause.isdigit():
                schema = self.mysql.get_table_schema(table_name)
                if schema["primary_keys"]:
                    pk = schema["primary_keys"][0]
                    query = f"SELECT * FROM {table_name} WHERE {pk} = {where_clause}"
                else:
                    query = f"SELECT * FROM {table_name} WHERE id = {where_clause}"
            else:
                query = f"SELECT * FROM {table_name} WHERE {where_clause}"
        
        success, result, error = self.mysql.execute_query(query)
        
        if not success:
            return ToolResult(success=False, data=None, error=f"查询失败：{error}")
        
        return ToolResult(
            success=True,
            data={
                "table": table_name,
                "records": result,
                "count": len(result) if result else 0
            },
            metadata={"operation": "get", "table": table_name, "query": query}
        )
    
    async def _set_operation(self, args: List[str]) -> ToolResult:
        if len(args) < 2:
            return ToolResult(success=False, data=None, error="Set 操作需要表名和 JSON 数据。")
        
        table_name = args[0]
        if not self.mysql.table_exists(table_name):
            return ToolResult(success=False, data=None, error=f"表 '{table_name}' 不存在")

        try:
            json_data = " ".join(args[1:])
            data = repair_json(json_data, return_objects=True)
            if not isinstance(data, dict):
                raise ValueError("数据必须是JSON格式")
        except (json.JSONDecodeError, ValueError) as e:
            return ToolResult(
                success=False,
                data=None,
                error=f"无效的JSON数据: {str(e)}.示例格式：set {table_name} {{\"column1\": \"value1\", \"column2\": \"value2\"}}.。键和字符串值都使用双引号。使用“mysql_database help”可获取更多示例。"
            )

        columns = list(data.keys())
        values = list(data.values())
        placeholders = ", ".join(["%s"] * len(values))
        query = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({placeholders})"
        success, result, error = self.mysql.execute_query(query, tuple(values))
        
        if not success:
            return ToolResult(success=False,data=None,error=f"数据插入失败：{error}")
        
        return ToolResult(
            success=True,
            data={
                "table": table_name,
                "inserted_data": data,
                "affected_rows": result
            },
            metadata={"operation": "insert", "table": table_name}
        )
    
    async def _update_operation(self, args: List[str]) -> ToolResult:
        if len(args) < 4:
            return ToolResult(success=False, data=None,
                error="更新操作需要: table, JSON data, 'where', and condition"
            )
        
        table_name = args[0]
        if not self.mysql.table_exists(table_name):
            return ToolResult(success=False, data=None, error=f"表 '{table_name}' 不存在")

        try:
            where_index = args.index("where")
            json_data = " ".join(args[1:where_index])
            where_clause = " ".join(args[where_index + 1:])
        except ValueError:
            return ToolResult(success=False, data=None, error="更新操作需要 'where' 条件")

        try:
            data = repair_json(json_data, return_objects=True)
            if not isinstance(data, dict):
                raise ValueError("数据必须是JSON格式")
        except (json.JSONDecodeError, ValueError) as e:
            return ToolResult(success=False, data=None, error=f"无效的JSON数据: {str(e)}")

        set_clauses = [f"{col} = %s" for col in data.keys()]
        query = f"UPDATE {table_name} SET {', '.join(set_clauses)} WHERE {where_clause}"
        success, result, error = self.mysql.execute_query(query, tuple(data.values()))
        
        if not success:
            return ToolResult(success=False,data=None,error=f"更新失败: {error}")
        
        return ToolResult(
            success=True,
            data={
                "table": table_name,
                "updated_data": data,
                "affected_rows": result,
                "where_clause": where_clause
            },
            metadata={"operation": "update", "table": table_name}
        )
    
    async def _delete_operation(self, args: List[str]) -> ToolResult:
        if len(args) < 3:
            return ToolResult(
                success=False,
                data=None,
                error="删除操作需要: table, 'where', and condition"
            )
        
        table_name = args[0]
        if not self.mysql.table_exists(table_name):
            return ToolResult(success=False, data=None, error=f"表 '{table_name}' 不存在")

        try:
            where_index = args.index("where")
            where_clause = " ".join(args[where_index + 1:])
        except ValueError:
            return ToolResult( success=False, data=None, error="为保证安全，删除操作需要 'where' 条件")

        query = f"DELETE FROM {table_name} WHERE {where_clause}"
        success, result, error = self.mysql.execute_query(query)
        
        if not success:
            return ToolResult(success=False, data=None, error=f"删除失败：{error}" )
        
        return ToolResult(
            success=True,
            data={
                "table": table_name,
                "affected_rows": result,
                "where_clause": where_clause
            },
            metadata={"operation": "delete", "table": table_name}
        )
    
    async def _search_operation(self, args: List[str]) -> ToolResult:
        if len(args) < 3:
            return ToolResult(
                success=False,
                data=None,
                error="搜索操作需要表名、列名和搜索值。使用方法：search <table> <column> <value>"
            )
        
        table_name = args[0]
        column_name = args[1]
        search_value = " ".join(args[2:])

        if not self.mysql.table_exists(table_name):
            return ToolResult(
                success=False,
                data=None,
                error=f"表 '{table_name}' 不存在. 可用表：{', '.join(self.mysql.get_table_names())}"
            )

        try:
            schema = self.mysql.get_table_schema(table_name)
            if column_name not in schema["columns"]:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"列“{column_name}”在表“{table_name}”中不存在。可用列：{', '.join(schema['columns'])}"
                )

            column_type = schema["column_types"].get(column_name, "").lower()

            if any(text_type in column_type for text_type in ['varchar', 'text', 'char']):
                query = f"SELECT * FROM {table_name} WHERE {column_name} LIKE %s"
                search_params = (f"%{search_value}%",)
                search_type = "partial_text_match"
            elif any(num_type in column_type for num_type in ['int', 'decimal', 'float', 'double']):
                try:
                    if '.' in search_value:
                        float(search_value)
                    else:
                        int(search_value)
                    query = f"SELECT * FROM {table_name} WHERE {column_name} = %s"
                    search_params = (search_value,)
                    search_type = "exact_numeric_match"
                except ValueError:
                    return ToolResult(
                        success=False,
                        data=None,
                        error=f"列 '{column_name}' 为数值型 ({column_type}) 但查询值 '{search_value}' 不是数值"
                    )
            elif any(date_type in column_type for date_type in ['date', 'time', 'timestamp']):
                query = f"SELECT * FROM {table_name} WHERE {column_name} = %s"
                search_params = (search_value,)
                search_type = "exact_date_match"
            else:
                query = f"SELECT * FROM {table_name} WHERE {column_name} = %s"
                search_params = (search_value,)
                search_type = "exact_match"

            success, result, error = self.mysql.execute_query(query, search_params)
            
            if not success:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"查询失败：{error}. Query: {query}"
                )
            
            return ToolResult(
                success=True,
                data={
                    "table": table_name,
                    "column": column_name,
                    "search_value": search_value,
                    "search_type": search_type,
                    "column_type": column_type,
                    "records": result,
                    "count": len(result) if result else 0
                },
                metadata={
                    "operation": "search",
                    "table": table_name,
                    "column": column_name,
                    "query": query
                }
            )
            
        except Exception as e:
            logger.error(f"搜索操作期间出错: {e}")
            return ToolResult(success=False, data=None, error=f"搜索操作失败: {str(e)}")
    
    async def _stats_operation(self) -> ToolResult:
        tables = self.mysql.get_table_names()
        table_stats = {}
        total_records = 0
        
        for table in tables:
            try:
                success, result, error = self.mysql.execute_query(f"SELECT COUNT(*) as count FROM {table}")
                if success and result:
                    count = result[0]['count']
                    table_stats[table] = {
                        "record_count": count,
                        "schema": self.mysql.get_table_schema(table)
                    }
                    total_records += count
            except Exception as e:
                logger.warning(f"获取表 {table} 的统计信息时出错: {e}")
                table_stats[table] = {"record_count": "Error", "schema": {}}
        
        stats = {
            "database": self.mysql.database,
            "host": self.mysql.host,
            "total_tables": len(tables),
            "total_records": total_records,
            "tables": tables,
            "table_stats": table_stats,
            "connection_status": "connected" if self.mysql.is_connected() else "disconnected"
        }
        
        return ToolResult(
            success=True,
            data=stats,
            metadata={"operation": "stats", "database": self.mysql.database}
        )
    
    async def _create_table_operation_improved(self, query: str) -> ToolResult:
        import re
        match = re.match(r'create_table\s+(\w+)\s*\((.+)\)', query.strip(), re.IGNORECASE)
        if not match:
            return ToolResult(
                success=False,
                data=None,
                error="无效的 create_table 格式。预期格式：create_table 表名 (列定义)"
            )
        
        table_name = match.group(1)
        column_definitions = match.group(2).strip()

        if self.mysql.table_exists(table_name):
            return ToolResult(success=False, data=None, error=f"表 '{table_name}' 已存在")
        
        sql_query = f"CREATE TABLE {table_name} ({column_definitions})"
        
        success, result, error = self.mysql.execute_query(sql_query)
        
        if not success:
            return ToolResult(
                success=False,
                data=None,
                error=f"创建表失败: {error}"
            )
        
        return ToolResult(
            success=True,
            data={
                "table_name": table_name,
                "created": True
            },
            metadata={
                "operation": "create_table",
                "table": table_name
            }
        )

    async def _create_table_operation(self, args: List[str]) -> ToolResult:
        if len(args) < 2:
            return ToolResult(
                success=False,
                data=None,
                error="创建表操作需要表名和列定义。"
            )
        
        table_name = args[0]
        column_definitions = " ".join(args[1:])

        if self.mysql.table_exists(table_name):
            return ToolResult(success=False, data=None, error=f"表 '{table_name}' 已经存在")
        
        query = f"CREATE TABLE {table_name} ({column_definitions})"
        
        success, result, error = self.mysql.execute_query(query)
        
        if not success:
            return ToolResult(
                success=False,
                data=None,
                error=f"Create table failed: {error}"
            )
        
        return ToolResult(
            success=True,
            data={
                "table_name": table_name,
                "column_definitions": column_definitions,
                "created": True
            },
            metadata={"operation": "create_table", "table": table_name}
        )
    
    async def _custom_query_operation(self, args: List[str]) -> ToolResult:
        if not args:
            return ToolResult(
                success=False,
                data=None,
                error="SQL 操作需要一个查询语句。使用方法：sql SELECT * FROM 表名"
            )
        
        query = " ".join(args).strip()

        dangerous_keywords = ['DROP', 'TRUNCATE', 'ALTER DATABASE', 'CREATE DATABASE']
        query_upper = query.upper()
        for keyword in dangerous_keywords:
            if keyword in query_upper:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"不允许执行危险操作“{keyword}”。仅允许执行 SELECT、INSERT、UPDATE、DELETE、SHOW、DESCRIBE 操作。"
                )
        
        try:
            success, result, error = self.mysql.execute_query(query)
            
            if not success:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"SQL 查询执行失败：{error}\n查询：{query}\n提示：检查您的 SQL 语法以及表/列名称。"
                )

            if isinstance(result, list):
                count = len(result)
                result_type = "records"
            else:
                count = result
                result_type = "affected_rows"
            
            return ToolResult(
                success=True,
                data={
                    "query": query,
                    "result": result,
                    "count": count,
                    "result_type": result_type
                },
                metadata={
                    "operation": "custom_query",
                    "query_type": query.split()[0].upper() if query else "UNKNOWN"
                }
            )
            
        except Exception as e:
            logger.error(f"自定义查询操作中出现意外错误：{e}")
            return ToolResult(
                success=False,
                data=None,
                error=f"执行 SQL 查询时出现意外错误：{str(e)}\n查询：{query}\n这可能是连接问题或内部错误。"
            )
    
    async def _help_operation(self) -> ToolResult:
        help_info = {
            "operations": {
                "list_tables": {
                    "description": "列出数据库中的所有表",
                    "usage": "list_tables",
                    "example": "list_tables"
                },
                "describe": {
                    "description": "显示表结构和列",
                    "usage": "describe <表名>",
                    "example": "describe users"
                },
                "get": {
                    "description": "从表中检索数据",
                    "usage": "get <表名> [条件]",
                    "examples": [
                        "get users",
                        "get users id=1",
                        "get users * 10"
                    ]
                },
                "set": {
                    "description": "向表中插入新记录",
                    "usage": "set <表名> <json数据>",
                    "example": 'set users {"name": "John", "email": "john@example.com", "age": 25}',
                    "important": "数据必须是带双引号的有效JSON格式"
                },
                "update": {
                    "description": "更新现有记录",
                    "usage": "update <表名> <json数据> WHERE <条件>",
                    "example": 'update users {"name": "Jane"} WHERE id=1'
                },
                "delete": {
                    "description": "从表中删除记录",
                    "usage": "delete <表名> WHERE <条件>",
                    "example": "delete users WHERE id=1"
                },
                "search": {
                    "description": "按列值搜索记录",
                    "usage": "search <表名> <列名> <值>",
                    "example": "search users name John"
                },
                "stats": {
                    "description": "显示数据库统计信息",
                    "usage": "stats",
                    "example": "stats"
                },
                "create_table": {
                    "description": "创建新表",
                    "usage": "create_table <表名> (<列定义>)",
                    "example": "create_table users (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100), age INT)"
                },
                "sql": {
                    "description": "执行自定义SQL查询",
                    "usage": "sql <查询语句>",
                    "example": "sql SELECT * FROM users WHERE age > 18"
                }
            },
            "common_errors": {
                "JSON parsing error": "在JSON中使用双引号：{\"key\": \"value\"}",
                "Table doesn't exist": "使用'list_tables'查看可用表",
                "Syntax error": "检查上述操作格式和示例"
            }
        }
        
        return ToolResult(
            success=True,
            data=help_info,
            metadata={"operation": "help"}
        )

    def __del__(self):
        if hasattr(self, 'mysql'):
            self.mysql.disconnect()

if __name__ == '__main__':
    import asyncio
    tool_case = MySQLDatabaseTool()
    print(tool_case.name)
    print(tool_case.description)
    print(tool_case.args)
    print(asyncio.run(tool_case.arun({
        'query': 'sql select * from users',
    })))